from __future__ import absolute_import

import torch
import torch.nn as nn
import torch.nn.functional as F
from model.utils.config import cfg
from model.loss.losses import _smooth_l1_loss
from torch.autograd import Variable

from .anchor_target_layer import _AnchorTargetLayer
from .proposal_layer import _ProposalLayer
from model.loss.losses import  hard_negative_mining




class _RPN(nn.Module):
    """ region proposal network """
    def __init__(self, din):
        super(_RPN, self).__init__()

        self.din = din  # get depth of input feature map, e.g., 512
        self.anchor_scales = cfg.ANCHOR_SCALES
        self.anchor_ratios = cfg.ANCHOR_RATIOS
        self.feat_stride = cfg.FEAT_STRIDE

        # define the convrelu layers processing input feature map
        # self.RPN_Conv = nn.Conv2d(self.din, 512, 3, 1, 1, bias=True)

        # define bg/fg classifcation score layer， 分类输出的通道数=anchor尺寸个数*长宽比个数*2（前景/背景）
        self.nc_score_out = len(self.anchor_scales) * len(
            self.anchor_ratios) * 2  # 2(bg/fg) * 9 (anchors)

        self.RPN_cls_score = nn.Conv2d(self.din, self.nc_score_out, 1, 1, 0)

        # define anchor box offset prediction layer，偏移预测的通道数=anchor尺寸个数*长宽比个数*4
        self.nc_bbox_out = len(self.anchor_scales) * len(
            self.anchor_ratios) * 4  # 4(coords) * 9 (anchors)
        self.RPN_bbox_pred = nn.Conv2d(self.din, self.nc_bbox_out, 1, 1, 0)

        # define proposal layer
        self.RPN_proposal = _ProposalLayer(self.feat_stride,
                                           self.anchor_scales,
                                           self.anchor_ratios)

        # define anchor target layer
        self.RPN_anchor_target = _AnchorTargetLayer(self.feat_stride,
                                                    self.anchor_scales,
                                                    self.anchor_ratios)

        self.rpn_loss_cls = 0
        self.rpn_loss_box = 0



    @staticmethod
    def reshape(x, d):
        input_shape = x.size()
        x = x.view(input_shape[0], int(d),
                   int(float(input_shape[1] * input_shape[2]) / float(d)),
                   input_shape[3])
        return x

    def forward(self, base_feat, im_info, gt_boxes, num_boxes):
        batch_size = base_feat.size(0)

        # return feature map after convrelu layer
        # rpn_conv1 = F.relu(self.RPN_Conv(base_feat), inplace=True)
        # get rpn classification score
        rpn_cls_score = self.RPN_cls_score(base_feat)# 512x512输入，feat_stride=16，(bs,9*2,32,32)

        rpn_cls_score_reshape = self.reshape(rpn_cls_score, 2)# reshape(bs, 2, 9*32, 32)
        rpn_cls_prob_reshape = F.softmax(rpn_cls_score_reshape, 1) # softmax(bs, 2, 9*32, 32)
        rpn_cls_prob = self.reshape(rpn_cls_prob_reshape, self.nc_score_out)# reshape(bs, 2*9, 32, 32)

        # get rpn offsets to the anchor boxes
        rpn_bbox_pred = self.RPN_bbox_pred(base_feat)# (bs,9*4,32,32)

        # proposal layer
        cfg_key = 'TRAIN' if self.training else 'TEST'

        rois = self.RPN_proposal(
            (rpn_cls_prob.data, rpn_bbox_pred.data, im_info, cfg_key))# 生成并筛选anchor。

        self.rpn_loss_cls = 0
        self.rpn_loss_box = 0

        # generating training labels and build the rpn loss
        if self.training:
            assert gt_boxes is not None

            rpn_data = self.RPN_anchor_target(
                (rpn_cls_score.data, gt_boxes, im_info, num_boxes))# anchor标定，返回labels分类标签，bbox_targe偏移标签，bbox_inside_weights,bbox_outside_weights

            # compute classification loss
            rpn_cls_score = rpn_cls_score_reshape.permute(
                0, 2, 3, 1).contiguous().view(batch_size, -1, 2)
            rpn_label = rpn_data[0].view(batch_size, -1)

            rpn_keep = Variable(rpn_label.view(-1).ne(-1).nonzero().view(-1))

            rpn_cls_score = torch.index_select(rpn_cls_score.view(-1, 2), 0,
                                               rpn_keep)

            rpn_label = torch.index_select(rpn_label.view(-1), 0,
                                           rpn_keep.data)
            rpn_label = Variable(rpn_label.long())

            # from collections import  Counter
            # label = rpn_label.cpu().numpy()
            # print(Counter(label))

            loss = -F.log_softmax(rpn_cls_score, dim=1)[:, 0]
            mask ,num_pos = hard_negative_mining(loss, rpn_label)
            confidence = rpn_cls_score[mask, :]
            self.rpn_loss_cls = F.cross_entropy(confidence.reshape(-1, 2), rpn_label[mask], reduction='mean')

            # self.rpn_loss_cls = F.cross_entropy(rpn_cls_score, rpn_label)
            # self.rpn_loss_cls = OHEM_loss(rpn_cls_score, rpn_label)




            fg_cnt = torch.sum(rpn_label.data.ne(0))

            rpn_bbox_targets, rpn_bbox_inside_weights, rpn_bbox_outside_weights = rpn_data[
                1:]

            # compute bbox regression loss
            rpn_bbox_inside_weights = Variable(rpn_bbox_inside_weights)
            rpn_bbox_outside_weights = Variable(rpn_bbox_outside_weights)
            rpn_bbox_targets = Variable(rpn_bbox_targets)




            self.rpn_loss_box = _smooth_l1_loss(rpn_bbox_pred,
                                                rpn_bbox_targets,
                                                rpn_bbox_inside_weights,
                                                rpn_bbox_outside_weights,
                                                sigma=3,
                                                dim=[1, 2, 3],
                                           )



        return rois, self.rpn_loss_cls, self.rpn_loss_box
